import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics
import numpy as np
import os
import sys
import pickle
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

# 指定随机种子
np.random.seed(777)
tf.random.set_seed(777)

# 超参数
ALPHA = 0.001
BATCH_SIZE = 64
N_EPOCH = 5  # 为了快速演示，这里设定偏小，正式代码请酌情加大
N_BLOCK_UNITS = 2  # inception block组数，每组2个block
N_INIT_CH = 32  # 初始通道数
print(f'alpha = {ALPHA}, batch_size = {BATCH_SIZE}, n_epochs = {N_EPOCH}')
print(f'n_block_units = {N_BLOCK_UNITS}， n_init_ch = {N_INIT_CH}')

# 1. 使用keras定义InceptionNet网络进行分类（每题6分，共60分）
# ①    数据处理
# 1)    加载cifar10数据集（6分）
with open('./data/cifar10/data_batch_1', 'br') as f:
    loaded = pickle.load(f, encoding='bytes')
# b'batch_label'
# ()
# b'labels'
# (10000,)
# b'data'
# (10000, 3072)
# b'filenames'
# (10000,)

# 2)    合理进行数据预处理（6分）
x = loaded[b'data']
x = x.astype(np.float32) / 255.0
x = x.reshape(-1, 3, 32, 32)  # (10000, 3, 32, 32)
x = x.transpose(0, 2, 3, 1)  # (10000, 32, 32, 3)
y = loaded[b'labels']
x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.9, random_state=777)

# 3)    将特征和标签数据放入通道中（6分）
ds_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))\
    .shuffle(buffer_size=1000)\
    .batch(BATCH_SIZE, drop_remainder=True)\
    .prefetch(tf.data.experimental.AUTOTUNE)
ds_test = tf.data.Dataset.from_tensor_slices((x_test, y_test))\
    .shuffle(buffer_size=1000)\
    .batch(BATCH_SIZE)\
    .prefetch(tf.data.experimental.AUTOTUNE)


# ②    模型搭建
# 1)    定义卷积类，在类内进行卷积，批量归一化，非线性转化(Relu)（6分）
class ConvBnRelu(keras.Model):

    def __init__(self, filters, kernel_size, strides, **kwargs):
        super().__init__(**kwargs)
        self.conv = layers.Conv2D(filters, kernel_size, strides, 'same');
        self.bn = layers.BatchNormalization()
        self.relu = layers.ReLU()

    def call(self, x, training=None):
        x = self.conv(x, training=training)
        x = self.bn(x, training=training)
        x = self.relu(x, training=training)
        return x


# 2)    参照下图完成InceptionNetBlack模块搭建（6分）
class InceptionNetBlock(keras.Model):

    def __init__(self, filters, strides, **kwargs):
        super().__init__(**kwargs)
        self.conv1 = ConvBnRelu(filters, (1, 1), strides)
        self.conv21 = ConvBnRelu(filters, (1, 1), strides)
        self.conv22 = ConvBnRelu(filters, (3, 3), (1, 1))
        self.conv31 = ConvBnRelu(filters, (1, 1), strides)
        self.conv32 = ConvBnRelu(filters * 2, (5, 5), (1, 1))
        self.pool41 = layers.MaxPool2D((3, 3), strides, 'same')
        self.conv42 = ConvBnRelu(filters * 2, (1, 1), (1, 1))

    # 3)    InceptionNetBlack模块进行正向传播（6分）
    def call(self, x, training=None):
        x1 = self.conv1(x, training=training)
        x2 = self.conv21(x, training=training)
        x2 = self.conv22(x2, training=training)
        x3 = self.conv31(x, training=training)
        x3 = self.conv32(x3, training=training)
        x4 = self.pool41(x, training=training)
        x4 = self.conv42(x4, training=training)
        r = tf.concat([x1, x2, x3, x4], 3)
        return r


# 4)    每使用两层InceptionNetBlack模块，通道数翻倍，图片尺寸减半，完成InceptionNet网络的构建（6分）
class InceptionNet(keras.Model):

    def __init__(self, init_ch, n_cls, n_block_units, **kwargs):
        super().__init__(**kwargs)
        filters = None
        self.blocks = keras.Sequential()
        for unit_id in range(n_block_units):
            for layer_id in range(2):
                if layer_id == 0:
                    strides = (2, 2)
                    if filters is None:
                        filters = init_ch
                    else:
                        filters *= 2
                else:
                    strides = (1, 1)
                block = InceptionNetBlock(filters, strides)
                self.blocks.add(block)
        self.gpool = layers.GlobalAvgPool2D()
        self.flt = layers.Flatten()
        self.fc1 = layers.Dense(200, activation=activations.relu)
        self.fc2 = layers.Dense(n_cls, activation=activations.softmax)

    # 5)    InceptionNet模块进行正向传播（6分）
    def call(self, x, training=None):
        x = self.blocks(x, training=training)
        x = self.gpool(x, training=training)
        x = self.flt(x, training=training)
        x = self.fc1(x, training=training)
        x = self.fc2(x, training=training)
        return x


# ③    模型预测
model = InceptionNet(N_INIT_CH, 10, N_BLOCK_UNITS)
model.build(input_shape=(None, 32, 32, 3))
model.summary()

# 1)    进行训练，选择Adam优化器，合理选择损失函数和迭代次数（6分）
model.compile(
    loss=losses.sparse_categorical_crossentropy,
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    metrics=[metrics.sparse_categorical_accuracy]
)
model.fit(ds_train,
          epochs=N_EPOCH,
          validation_data=ds_test)

# 2)    计算并打印测试集的准确率（6分）
print('计算并打印测试集的准确率')
model.evaluate(ds_test)
